import csv
import glob
import os
from collections import defaultdict

def find_highest_avg_speedup_by_batch():
    """找出每个 batch size 下最高的平均 speedup"""
    
    print("=== 每个 Batch Size 下最高的平均 Speedup ===\n")
    
    # 获取所有带 speedup 的文件
    speedup_files = glob.glob("timings_all_setting_rlhf_*_with_speedup.csv")
    
    # 按 batch_size 收集数据
    batch_file_avg_speedups = defaultdict(lambda: defaultdict(list))
    
    for file_path in sorted(speedup_files):
        file_name = os.path.basename(file_path)
        
        with open(file_path, newline='') as f:
            reader = csv.DictReader(f)
            for row in reader:
                if row['speedup'] and row['speedup'] != '':
                    speedup = float(row['speedup'])
                    batch_size = int(row['batch_size'])
                    batch_file_avg_speedups[batch_size][file_name].append(speedup)
    
    # 计算每个 batch_size 下每个文件的平均 speedup
    batch_file_avg = defaultdict(dict)
    for batch_size in batch_file_avg_speedups:
        for file_name, speedups in batch_file_avg_speedups[batch_size].items():
            avg_speedup = sum(speedups) / len(speedups)
            batch_file_avg[batch_size][file_name] = avg_speedup
    
    # 按 batch_size 输出最高平均 speedup
    print("batch_size | 最高平均speedup | 文件")
    print("-" * 50)
    for batch_size in sorted(batch_file_avg.keys()):
        max_avg_speedup = 0
        best_file = ""
        
        for file_name, avg_speedup in batch_file_avg[batch_size].items():
            if avg_speedup > max_avg_speedup:
                max_avg_speedup = avg_speedup
                best_file = file_name
        
        print(f"{batch_size:9d} | {max_avg_speedup:13.4f} | {best_file}")
    
    # 额外统计：显示所有文件的平均 speedup
    print("\n=== 每个 Batch Size 下所有文件的平均 Speedup ===\n")
    for batch_size in sorted(batch_file_avg.keys()):
        print(f"Batch Size {batch_size}:")
        print("-" * 40)
        
        # 按平均 speedup 排序
        sorted_files = sorted(batch_file_avg[batch_size].items(), key=lambda x: x[1], reverse=True)
        
        for file_name, avg_speedup in sorted_files:
            data_count = len(batch_file_avg_speedups[batch_size][file_name])
            print(f"  {file_name}: 平均 {avg_speedup:.4f}, 数据点 {data_count}")
        
        print()

if __name__ == "__main__":
    find_highest_avg_speedup_by_batch() 